import torch
import torch.nn as nn

import math

class TransMitModel:
    def __init__(self, Nc, Q):
        self.Nc = Nc
        self.Q = Q
    
    def transmit(self, signal, noise, H):
        z = H @ signal
        z = z + noise  # [B, Nc, K, 2*Q]
        z = z.permute(0,2,1,3) # [B, K, Nc, 2*Q] 复数
        z = torch.cat((torch.real(z),torch.imag(z)), 3) # [B, K, Nc, 2*Q]
        z = z.reshape(-1, self.Nc, 2*self.Q) # [B*K, Nc, 2*Q]
        return z 